import requests
from requests_toolbelt import MultipartEncoder
import os
import time
import json
import webbrowser
import contextlib
import mimetypes
from dotenv import load_dotenv

load_dotenv()

API_URL = "https://ai.gitee.com/v1/async/documents/parse"
api_key = os.getenv("api_key")
headers = {
    "Authorization": f"Bearer {api_key}" # 使用 Bearer 令牌进行身份验证
}

def query(payload):
    fields = [
        ("model", payload["model"]),
        ("is_ocr", str(payload["is_ocr"]).lower()),
        ("formula_enable", str(payload["formula_enable"]).lower()),
        ("table_enable", str(payload["table_enable"]).lower()),
        ("layout_model", payload["layout_model"]),
    ]
    with contextlib.ExitStack() as stack:
        filepath = payload["file"]
        name = os.path.basename(filepath)
        if filepath.startswith(("http://", "https://")):
            response = requests.get(filepath, timeout=10)
            response.raise_for_status()
            fields.append(("file", (name, response.content, response.headers.get("Content-Type", "application/octet-stream"))))
        else:
            mime_type, _ = mimetypes.guess_type(filepath)
            fields.append(("file", (name, stack.enter_context(open(filepath, "rb")), mime_type or "application/octet-stream")))
        encoder = MultipartEncoder(fields)
        headers["Content-Type"] = encoder.content_type
        response = requests.post(API_URL, headers=headers, data=encoder)
        return response.json()

def poll_task(task_id):
    status_url = f"https://ai.gitee.com/v1/task/{task_id}"
    timeout = 30 * 60
    retry_interval = 10
    attempts = 0
    max_attempts = int(timeout / retry_interval)
    while attempts < max_attempts:
        attempts += 1
        print(f"Checking task status [{attempts}]...", end="")
        response = requests.get(status_url, headers=headers, timeout=10)
        result = response.json()
        if result.get("error"):
            print('error')
            raise ValueError(f"{result['error']}: {result.get('message', 'Unknown error')}")
        status = result.get("status", "unknown")
        print(status)
        if status == "success":
            output = result.get("output", {})
            segments = output.get("segments", [])
            if segments:
                print("\n===== Recognition Results =====")
                for seg in segments:
                    content = seg.get("content", "")
                    print(content)
            else:
                print("⚠️ No content found in output")
            task_file = f"task_{task_id}.json"
            with open(task_file, "w") as f:
                json.dump(result, f, indent=4)
            print(f"Task was saved to file {task_file}")
            return result
        elif status in ["failed", "cancelled"]:
            print(f"❌ Task {status}")
        else:
            time.sleep(retry_interval)
            continue
    print(f"⏰ Maximum attempts reached ({max_attempts})")
    return {"status": "timeout", "message": "maximum wait time exceeded"}

if __name__ == "__main__":
    print("Creating task...")
    result = query({
        "file": "path/to/file.pdf", # 文件路径，用户需要替换为实际的文件路径
        "model": "PDF-Extract-Kit-1.0",
        "is_ocr": True,
        "formula_enable": True,
        "table_enable": True,
        "layout_model": "doclayout_yolo"
    })
    task_id = result.get("task_id")
    if not task_id:
        raise ValueError("Task ID not found in the response")
    print(f"Task ID: {task_id}")
    task = poll_task(task_id)
    if task.get("status") == "success":
        print("Task completed successfully!")